from prefect import task
import pandas as pd
from prefect import get_run_logger

from application.db.redis_db.base_redis import RedisManager
from application.tasks.base_task import BaseTask
from application.utils.utils import translation_text_content

redis_manager = RedisManager()
redis_conn = redis_manager.get_connection()


# BaseTask
class ArticleChFieldTranslationTask(BaseTask):
    def run(self, df: pd.DataFrame):
        logger = get_run_logger()
        logger.info("Computing statistics")
        check_columns = [
            'check_article_info_does_it_exist',
            'check_whether_the_language_type_is_wrong',
            'check_chinese_and_english_titles_are_missing',
            'check_published_time_precision',
            'check_whether_the_journal_data_exists'
        ]
        # 只保留article_id为空的数据
        # df = df[df['article_id'].isna()]
        df = df[df[check_columns].all(axis=1)]
        # 检查article_keywords列是否存在，如果不存在则创建空列
        if 'article_keywords' not in df.columns:
            df['article_keywords'] = None
        if 'article_abstract' not in df.columns:
            df['article_abstract'] = None

        df['article_info'] = df['article_info'].apply(self.article_info_translate)  # 文献翻译
        df.loc[:, 'article_abstract'] = df['article_abstract'].apply(self.article_abstract_translate)
        df.loc[:, 'article_keywords'] = df['article_keywords'].apply(self.article_keywords_translate)

        return df

    def article_info_translate(self, article_info):
        """
        文献类型翻译
        :param article_info:
        :return:
        """
        type_name = article_info.get("type_name") or {}
        article_name = article_info.get("article_name") or {}

        if not type_name.get("chi") and type_name.get("eng"):
            # 译文翻译
            chs_text = translation_text_content(type_name.get("eng") or "")
            if not chs_text['result']:
                return {"result": False, "msg": chs_text['msg']}
            chs_text = chs_text['data']
            print(f"已翻译文献类型:【{type_name.get('eng')}】->【{chs_text}】")
            type_name['chi'] = chs_text
            type_name['type_translation'] = True
        else:
            type_name['type_translation'] = False

        if not article_name.get("chi"):
            # 译文翻译
            chs_text = translation_text_content(article_name.get("eng") or "")
            if not chs_text['result']:
                return {"result": False, "msg": chs_text['msg']}
            chs_text = chs_text['data']
            print(f"已翻译标题:【{article_name.get('eng')}】->【{chs_text}】")

            article_name['chi'] = chs_text
            article_name['type_translation'] = True
        else:
            article_name['type_translation'] = False

        return article_info

    def article_abstract_translate(self, article_abstract):
        """
        文献摘要翻译
        :param article_abstract:
        :return:
        """
        # 确认摘要段落是否需要翻译
        if not article_abstract.get("chs_info"):
            for item in (article_abstract.get("eng_info") or []):
                # 译文翻译
                chs_text = translation_text_content(item.get("section_text") or "")
                if not chs_text['result']:
                    return {"result": False, "msg": chs_text['msg']}
                chs_text = chs_text['data']
                print(f"已翻译摘要段落:【{item.get("section_text")}】->【{chs_text}】")

                # 更新译文
                article_abstract.setdefault("chs_info", []).append({
                    "section_attr": item.get("section_attr"),
                    "section_text": chs_text,
                })
        # 对摘要剔空
        article_abstract = {k: v for k, v in article_abstract.items() if v}
        return article_abstract

    def article_keywords_translate(self, article_keywords):
        """
        文献关键词翻译
        :param article_keywords:
        :return:
        """
        if not isinstance(article_keywords, list):
            return article_keywords
        print(f"开始翻译关键词{article_keywords}")
        for index, item in enumerate(article_keywords):
            if not item.get("keyword_chs"):
                # 译文翻译
                keyword_chs = translation_text_content(item.get("keyword_eng") or "")
                if not keyword_chs['result']:
                    return {"result": False, "msg": keyword_chs['msg']}
                keyword_chs = keyword_chs['data']
                if len(keyword_chs) > len(item.get("keyword_eng") or "") * 3:
                    keyword_chs = item.get("keyword_eng") or ""
                print(f"已翻译关键词:【{item.get('keyword_eng')}】->【{keyword_chs}】")
                # 更新译文
                article_keywords[index]['keyword_chs'] = keyword_chs
            if not item.get("keyword_eng"):
                article_keywords[index]['keyword_eng'] = ''
        return article_keywords